# %% [markdown]
# ## ode solver

# %%
import numpy as np
from scipy.integrate import ode

#********************* ode 物理参数
mass: float = 1.0  # 质量或转动惯量
beta: float = 1.0  # ODE 阻尼系数

#********************* 质量矩阵
massM0 = np.array([[1.0, 0.0], [0.0, mass]], dtype=float)
inv_massM = np.linalg.inv(massM0)


#********************* 外力 或者 力矩;
# 对应 [坐标系平移, 加速度], 例如 [0.0,0.2] 对应外力为 0.2
def force(t):
    return np.array([0.0, 0.2 * np.sin(t)], dtype=float)
    return [0.0, 0.2]


#********************* jacobian 矩阵 df/dy; ODE beta 项
betaM = np.array([[0, 1], [0, -beta]], dtype=float)

#********************* ode 初始 [位移,速度], 初始时间
y0, t0 = [1.0, 0.0], 0.0


#********************* y'(t) = f(t, y); jac = df/dy
def f(t, y, arg1):
    # print(f'betaM is {betaM}')
    # print(y)
    tmp = betaM @ y + force(t)
    return inv_massM @ tmp


def jac(t, y, arg1):
    return betaM


print(betaM)
print(inv_massM)

# %%
# ********************* ODE参数, 终止时间, 步长
t1: float = 12.0  # 终止时刻
dt: float = 0.1  # 步长
r = ode(f, jac).set_integrator('vode', method='bdf', order=5)
r.set_initial_value(y0, t0).set_f_params(0.0).set_jac_params(0.0)
pre_step_count = int(t1 / dt) + 3  #
# debug test
if 0:
    print(r.successful())
    print(r.t)
    print(r.integrate(r.t + dt))

np.set_printoptions(formatter={'float': '{:< #9.8g}'.format})
#
xvArray = np.zeros((pre_step_count, 3), dtype=float)
xvArray[0] = [0.0, *y0]
ode_cnt: int = 1
while r.successful() and r.t < t1:
    t11 = r.t + dt
    ret = r.integrate(r.t + dt)
    xvArray[ode_cnt] = [t11, *ret]
    print(f'{t11:<#5.3g}, {ret}')
    ode_cnt += 1

# %% 画图
from matplotlib import pyplot as plt

fig, axes = plt.subplots(1, 2)
line1, = axes[0].plot(xvArray[:ode_cnt, 0],
                      xvArray[:ode_cnt, 1],
                      'b-',
                      label='x')
axes[0].legend(handles=[line1])
line2, = axes[1].plot(xvArray[:ode_cnt, 0],
                      xvArray[:ode_cnt, 2],
                      'g-',
                      label='v')
axes[1].legend(handles=[line2])
plt.show()
